import numpy as np
from sklearn.metrics import classification_report, roc_curve

from ml_rs import data_read
from ml_rs.vis_function import get_cls_func, roc_plot


def read_data():
    global train_x, train_y, test_x, test_y
    [train_x, train_y] = data_read.read_in("..\\train_data")
    [test_x, test_y] = data_read.read_in("..\\test_data")


def classification():
    global cla_num
    cla_num = input("Please enter a category: (1) II Classification (2) V Classification")
    if cla_num == "1":
        train_y[train_y == "train_roof"] = "0"
        train_y[train_y == "train_path"] = "0"
        train_y[train_y == "train_road"] = "0"
        train_y[train_y == "train_grass"] = "1"
        train_y[train_y == "train_tree"] = "1"

        test_y[test_y == "test_roof"] = "0"
        test_y[test_y == "test_path"] = "0"
        test_y[test_y == "test_road"] = "0"
        test_y[test_y == "test_grass"] = "1"
        test_y[test_y == "test_tree"] = "1"
    else:
        train_y[train_y == "train_roof"] = "roof"
        train_y[train_y == "train_path"] = "path"
        train_y[train_y == "train_road"] = "road"
        train_y[train_y == "train_grass"] = "grass"
        train_y[train_y == "train_tree"] = "tree"

        test_y[test_y == "test_roof"] = "roof"
        test_y[test_y == "test_path"] = "path"
        test_y[test_y == "test_road"] = "road"
        test_y[test_y == "test_grass"] = "grass"
        test_y[test_y == "test_tree"] = "tree"


def create_model():
    # 创建模型
    name = input('classification method:（1）LogisticRegression （2）SVM  （3）LDA')
    clf = get_cls_func(name)  # 选取分类器
    global model
    model = clf.fit(train_x, train_y)  # 训练分类器
    # 特征系数
    print("特征系数：", model.coef_)
    # 截距
    print("截距：", model.intercept_)


def predict_model():
    print("模型预测：")
    global predict_y, proba
    predict_y = model.predict(test_x)
    print(model.predict(test_x))  # 数据的类别
    proba = model.predict_proba(test_x)
    print(model.predict_proba(test_x))  # 数据的概率


# 精度报告
def clf_report():
    if cla_num == "1":
        cla_name = ['0', '1']
    else:
        cla_name = ["roof", "path", "road", "grass", "tree"]
    print('精度报告：')
    print(classification_report(test_y, predict_y, target_names=cla_name))


def draw_roc():
    if cla_num == "1":
        test_y[test_y == '1'] = 1
        test_y[test_y == '0'] = 0
        roc_plot(test_y.tolist(), proba[:, 1])


if __name__ == '__main__':
    read_data()
    classification()
    create_model()
    predict_model()
    clf_report()
    draw_roc()
